# -*-coding: utf-8-*-

import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from pylab import mpl


mpl.rcParams['font.sans-serif'] = ['FangSong'] # 指定默认字体
mpl.rcParams['axes.unicode_minus'] = False # 解决保存图像是负号'-'显示为方块的问题


def divid(source:pd.DataFrame, result:pd.DataFrame):
    print(source)
    isin = source['_id'].astype(result['_id'].dtype).isin(result['_id'])
    navy = source[isin]
    not_navy = source[-isin]
    return (navy,not_navy)


def prepare(data: pd.DataFrame, x, y, z):
    data[[x, y, z]] = data[[x, y, z]].apply(lambda x: np.log(x+1))
    data[x] = data[x] - pd.np.mean(data[x])
    data[y] = data[y] - pd.np.mean(data[y])
    data[z] = data[z] - pd.np.mean(data[z])
    return data


class DrawColor(object):
    @property
    def raw_color(self):
        return 'b'

    @property
    def other_color(self):
        return 'y'

    @property 
    def navy_color(self):
        return 'r'


class ClusterDrawResult(DrawColor):
    def __init__(self, source_file, result_file, 
            x='n_fans', y='n_follows', z='tweets'):
        self.X = x
        self.Y = y
        self.Z = z
        self.meta = prepare(pd.read_csv(source_file, 
                usecols=['_id', x, y, z]), x, y, z)
        result = pd.read_csv(result_file, usecols=['_id'])
        navy, not_navy = divid(self.meta, result)
        self.navy = navy
        self.not_navy = not_navy
        self.paint = Axes3D(plt.figure())

    def draw_raw(self):
        self.paint.scatter(self.meta[self.X],
                self.meta[self.Y],
                self.meta[self.Z], s=1, c=self.raw_color)
        self.paint.legend(['原始数据'])
        plt.show()
    
    def draw_rsl(self):
        self.paint.scatter(
                self.navy[self.X],
                self.navy[self.Y],
                self.navy[self.Z],
                s=5,
                c=self.navy_color)
        self.paint.scatter(
                self.not_navy[self.X],
                self.not_navy[self.Y],
                self.not_navy[self.Z],
                s=1,
                c=self.other_color)
        self.paint.legend(['水军', '非水军'])
        plt.show()

    def __del__(self):
        plt.close()


class SvmDraw(DrawColor):
    def __init__(self, sourcefile, x='n_fans', y='n_follows', z='tweets'):
        self.X = x
        self.Y = y
        self.Z = z
        self.data = prepare(pd.read_csv(
                                sourcefile, 
                                usecols=[x, y, z, 'label']),
                            x, y, z)
        self.navy, self.not_navy = self.divid(self.data)
        self.ax1 = Axes3D(plt.figure(1))
        self.ax2 = Axes3D(plt.figure(2))

    def divid(self, data):
        navy = data[data['label'] == 1]
        not_navy = data[data['label'] == 0]
        return navy, not_navy

    def draw(self):
        self.ax1.scatter(
            self.data[self.X],
            self.data[self.Y],
            self.data[self.Z],
            s=10,
            c=self.raw_color)
        self.ax2.scatter(
            self.navy[self.X],
            self.navy[self.Y],
            self.navy[self.Z],
            s=10,
            c=self.navy_color)
        self.ax2.scatter(
            self.not_navy[self.X],
            self.not_navy[self.Y],
            self.not_navy[self.Z],
            s=10,
            c=self.other_color)
        self.ax1.legend(['原始数据'])
        self.ax2.legend(['水军', '非水军'])
        plt.show()

    def __del__(self):
        plt.close()


if __name__ == '__main__':
    # ClusterDrawResult('meta.csv', 'result.csv').draw_raw()
    SvmDraw('svm_chuli.csv').draw()